from roboticstoolbox import DHRobot, RevoluteDH
import roboticstoolbox as rtb
import numpy as np
from spatialmath import SE3
import json
import time
from RRTSTAR import rrt_star_6d

class Jaka(DHRobot):
    def __init__(self):
        super().__init__(
            [
                RevoluteDH(alpha=np.pi/2, d=120.15, name="Joint 1"),
                RevoluteDH(a=-430, name="Joint 2"),
                RevoluteDH(a=-368.5, name="Joint 3"),
                RevoluteDH(alpha=np.pi/2, d=113.5, name="Joint 4"),
                RevoluteDH(alpha=-np.pi/2, d=113.5, name="Joint 5"),
                RevoluteDH(d=107, name="Joint 6")
            ],
            name="Jaka Robot"
        )


def get_object_pose(sim):
    # while True:     # 检查状态文件是否存在
    #     with open("status.json", "r") as f:
    #         data = json.load(f)
    #     if data["is_completed"]:
    #         pose_matrix = data["pose_matrix"]
    #         print("程序已完成，位姿矩阵为：", pose_matrix)
    #         # 在这里继续后续程序逻辑
    #         break
    #     else:
    #         print("程序尚未完成，继续检查...")
    #         time.sleep(1)

    ###存放代码####################################
    # pose_matrix = np.array([
    # [1, 0, 0, 1],
    # [0, 1, 0, 2],
    # [0, 0, 1, 3],
    # [0, 0, 0, 1]
    # ]).tolist()

    # data = {
    #     "is_completed": False,
    #     "pose_matrix": pose_matrix
    # }

    # with open("status.json", "w") as f:
    #     json.dump(data, f)
    ###存放代码####################################


    matrix =sim.getObjectMatrix(sim.getObject("/Cuboid"))
    matrix2 =sim.getObjectMatrix(sim.getObject("/hand_respondable"))
    
    T_sim_obj = np.array([
        [matrix[0], matrix[1], matrix[2], matrix[3]*1000],
        [matrix[4], matrix[5], matrix[6], matrix[7]*1000],
        [matrix[8], matrix[9], matrix[10], matrix[11]*1000],
        [0, 0, 0, 1]
    ])


    Hand_T= np.array([
        [matrix2[0], matrix2[1], matrix2[2], matrix2[3]*1000],
        [matrix2[4], matrix2[5], matrix2[6], matrix2[7]*1000],
        [matrix2[8], matrix2[9], matrix2[10], matrix2[11]*1000],
        [0, 0, 0, 1]
    ])
    print('T_sim_obj:', T_sim_obj,flush=True)
    print('Hand_T:', Hand_T,flush=True)
    return SE3(T_sim_obj)

def get_staight_trajectory(T_start,T_end,robot,start_q=[0,0,0,0,0,0]):
    cartisanean_trajectory = rtb.ctraj(T_start, T_end, 10)
    viapoints = []
    last_q = start_q
    for index, matrix in enumerate(cartisanean_trajectory):
        q= robot.ikine_LM(matrix,q0=last_q)
        last_q = q.q 
        viapoints.append(q.q)  
    viapoints = np.array(viapoints)
    trajectory = rtb.mstraj(viapoints , dt=0.1, tacc=0.2, qdmax=5)
    return trajectory


def get_trajectory(T_sim_obj,intial_angles):
    robot = Jaka()
    T_plan_end_initial = robot.fkine(robot.q)
    T_sim_end_initial = SE3(np.array([
        [1, 0, 0, -898.5],
        [0, 0, -1, -221],
        [0, 1, 0, 608.2],
        [0, 0, 0, 1]
    ]))
    T_sim_to_plan = T_plan_end_initial * T_sim_end_initial.inv()

    T_obj_to_precatch = SE3(np.array([
        [1, 0, 0, 1.5],
        [0, -1, 0, -246.1],
        [0, 0, -1, 583.2],
        [0, 0, 0, 1]
    ]))
    T_sim_end_target = T_obj_to_precatch * T_sim_obj
    T_sim_initial_to_target = T_sim_end_target * T_sim_end_initial.inv()
    T_plan_end_target = T_sim_to_plan * T_sim_initial_to_target * T_sim_to_plan.inv() * T_plan_end_initial

    T_plan_end_target = SE3.Trans(0,0,50)*T_plan_end_target


    T_plan_end_target.t[1]=-T_plan_end_target.t[1]-250

    print('T_plan_end_target:', T_plan_end_target,flush=True)
    print('TRANSLATE:', T_plan_end_target.t,flush=True)
    

    qf = robot.ikine_LM(T_plan_end_target, q0=intial_angles)

    viapoints = rrt_star_6d(intial_angles, qf.q)
    viapoints = np.array(viapoints)
    
    trajectory = rtb.mstraj(viapoints, dt=0.1, tacc=0.2, qdmax=5)
    # trajectory = rtb.jtraj(intial_angles, qf.q, t=50)
    



    T_catch_end_target = SE3.Trans(0,0,-150)*T_plan_end_target

    trajectory1 = get_staight_trajectory(T_plan_end_target,T_catch_end_target,robot,qf.q)

    T_move1 = SE3.Trans(0,0,200)*T_catch_end_target

    trajectory2 = get_staight_trajectory(T_catch_end_target,T_move1,robot,trajectory1.q[-1])

    T_move2 = SE3.Rx(-180, unit='deg')*T_move1
    qf = robot.ik_GN(T_move2)
    trajectory3 = rtb.jtraj(trajectory2.q[-1], qf[0], t=20)

    return trajectory , trajectory1 , trajectory2 , trajectory3


